import os
import tempfile

from charmicat import BadCommandLineError
from batchcat.jobscripts import SlurmScript
from csbbelle import BASFScript

class BASFSlurmScript(SlurmScript):
    script_body = '''#!/bin/bash
#SBATCH -A belle
#SBATCH -t {bcat_duration}
#SBATCH -N {bcat_nnodes}
#SBATCH -o {bcat_stdout}
#SBATCH -e {bcat_stderr}
#SBATCH -J {bcat_job_name}

source $HOME/.brofile

source /pic/projects/belle/scripts/general_env.sh
source $HOME/belle/belle_custom_setup.sh

shopt -s expand_aliases
{belleset}

{module_path}
echo CRAIGB_BASF_START `date`
basf <<EOF
{basf}
EOF

BASFRET=$?
echo CRAIGB_BASF_FINISH `date`

{stage}

if [ $BASFRET -ne 0 ]; then
  for p in {final_stub}*.h5; do
    mv $p $p.badret
  done
fi

'''
    
    stage_dir = tempfile.gettempdir()

    def __init__(self, name, module_name, input_file, opts):
        super(BASFSlurmScript, self).__init__(name, module_name, opts)

        input_file = os.path.abspath(input_file)

        if opts.bcat_outdir is None:
            opts.bcat_outdir = os.getcwd()

        if opts.stage:
            outpath_stub = os.path.join(BASFSlurmScript.stage_dir, 
                                        os.path.basename(input_file)\
                                            .rsplit('.', 1)[0])

            opts.stage = 'cp -v {outpath_stub}* {bcat_outdir}'
            final_stub = os.path.join(opts.bcat_outdir, 
                                      os.path.basename(outpath_stub))
            
        else:
            opts.stage = ''
            outpath_stub = os.path.join(opts.bcat_outdir, 
                                        os.path.basename(input_file)\
                                            .rsplit('.', 1)[0])
            final_stub = outpath_stub
            
        if opts.module_path is None:
            opts.module_path = ''
        else:
            opts.module_path = "export LD_LIBRARY_PATH={0}:$LD_LIBRARY_PATH"\
                .format(opts.module_path)

        if opts.basf_path is None:
            if 'outpath_stub' not in opts.basf_pars:
                opts.basf_pars['outpath_stub'] = outpath_stub
            
            if 'inpath' not in opts.basf_pars:
                opts.basf_pars['inpath'] = input_file

            if opts.hbook:
                basf_obj = BASFScript(module_name, input_file, opts.basf_pars, 
                                      outpath_stub + '.hbook')
            else:
                basf_obj = BASFScript(module_name, input_file, opts.basf_pars)

            basf = str(basf_obj)

        else:
            with open(opts.basf_path) as fp:
                basf = fp.read()

        if opts.bcat_job_name is None:
            opts.bcat_job_name = name

        if opts.hbook:
            belleset='belleset_hbook32768_libs'
        if opts.nocrash:
            belleset += '\n' + 'belleset_nocrash_libs'

        all_vars = {} 
        all_vars.update(locals())
        all_vars.update(vars(opts))
                
        self.script = BASFSlurmScript.script_body.format(**all_vars)

    def __str__(self):
        return self.script

            
